"""BCN类
它本身不存储结构，共享结构。只尺寸B位和N位
实现了单替换、禁邻位单替换、等双替换、限环等替换、不等替换的生成
"""
import os
# os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from polyhedron import Polyhedron
from cmds import Cluster, read, write_car, Periodic_Table
from copy import deepcopy
from random import sample, choice
import itertools
from functools import reduce
from itertools import combinations
import numpy as np
from C60_group_element import C60Ih_group, C60Ih轨道


class BCN(Polyhedron, Cluster):
    cluster = read(r'C60.car')
    framework = Polyhedron.from_cluster(cluster)
    g = None
    散列表 = set()

    def __init__(self, boron_pos=None, nitride_pos=None):
        self.B_pos = [] if boron_pos is None else boron_pos
        self.N_pos = [] if nitride_pos is None else nitride_pos

    def elements(self):
        elem_list = ['C'] * len(BCN.cluster)
        for p in self.B_pos:
            elem_list[p] = 'B'
        for p in self.N_pos:
            elem_list[p] = 'N'
        return elem_list

    def get_name(self):
        if not self.B_pos:
            name = 'C%dN%d'%(len(BCN.cluster)-len(self.N_pos), len(self.N_pos))
        elif not self.N_pos:
            name = 'C%dB%d'%(len(BCN.cluster)-len(self.B_pos), len(self.B_pos))
        else:
            name = 'C%dB%dN%d'%(len(BCN.cluster)-len(self.B_pos)-len(self.N_pos), len(self.B_pos), len(self.N_pos))
        return name+reduce(str.__add__, ['_%02d' %_ for _ in self.B_pos + self.N_pos])

    def __init_vertex_hash(self):
        return [Periodic_Table.element_id(e) for e in self.elements()]

    def __hash__(self):
        '''聚合一次，跟镜面hash取较大的'''
        vh = self.__init_vertex_hash()
        vh2a = self.framework._hash_aggregate3(vh)

        m = self.framework.mirror()
        vh2b = m._hash_aggregate3(vh)
        return max(sum(vh2a), sum(vh2b)) & 0xffffffffffffffff

    def adjacent(self):
        pass

    def save(self, folder_name):
        #file_name = os.path.join(folder_name, '%016x.car'%hash(self))
        file_name = os.path.join(folder_name, self.get_name()+'.car')
        c = self.cluster.deepcopy()
        c.set_element(self.elements())
        write_car(c, file_name)
        
    @staticmethod
    def 骨架ring():     
        rings, rings5, rings6 = [], [], []
        _,edge2face = BCN.framework.dual()
        calc_v = set() #已计算过的对偶点
        for e,v in edge2face.items():
            if v in calc_v:
                continue
            calc_v.add(v)
            ring = BCN.framework.get_face(e[0],e[1])
            rings.append(set(ring))
            if len(ring) == 5:
                rings5.append(set(ring))
            if len(ring) == 6:
                rings6.append(set(ring))
        return rings, rings5, rings6

    @staticmethod
    def 替换散列(a):
        '''一个替换的散列，用60进制表示替换位'''
        return reduce(lambda x, y: x * len(BCN.cluster) + y, [int(_) for _ in a])  #强制转int型，否则numpy.int跟int混用会出错！！！

    @staticmethod
    def 最小等价替换散列(B位, N位):
        等价替换 = C60Ih轨道(B位+N位, BCN.g)
        k = len(B位)
        return min([BCN.替换散列(sorted(_[:k])+sorted(_[k:])) for _ in 等价替换])
        
    @staticmethod
    def 重复否(替换):
        等价替换 = C60Ih轨道(替换, BCN.g)
        h = BCN.替换散列(替换)
        for rep in 等价替换:
            if BCN.替换散列(rep) < h:
                return True
        else:
            return False

    @staticmethod
    def 生成初始化():
        BCN.g = C60Ih_group()
        BCN.散列表 = set()
        BCN.计数 = 0
        k = 3
        #BCN.save_folder = r'C:\BCN\pre_data\C%dN%d_forbidden'%(len(BCN.cluster)-k,k)
        p,q = 4,4
        #BCN.save_folder = 'D:\\cluster\\C%dB%dN%d_ring2'%(60-p-q,p,q)
        #if not os.path.exists(BCN.save_folder):
        #    os.mkdir(BCN.save_folder)

    @staticmethod
    def 找到结构(B位, N位):
        """判断该结构是否重复，不重复则处理该新找到结构"""
        #限定BN键个数不低于某个值
        '''bonds = 0
        for i in B位:
            for j in N位:
                if j in BCN.framework.at[i]:
                    bonds += 1
        if bonds <= 2*len(B位):
            return'''

        min_h = BCN.最小等价替换散列(B位, N位)
        if min_h in BCN.散列表:
            return False
        BCN.散列表.add(min_h)

        BCN.计数 += 1
        if BCN.计数 % 100000 == 0:
            print(B位, N位, BCN.计数)

        bcn = BCN(N位, B位)
        f = open(r'D:\360disk\Cluster\topology\BCN\C49N11_五间.txt', 'a').write(bcn.get_name()+'\n') #仅保存字符串
        #bcn.save(BCN.save_folder) #保存结构文件
        return True

    @staticmethod
    def 单替换(k):
        BCN.生成初始化()
        for p in combinations(range(1,60), k-1):
            p1 = (0,)+p
            BCN.找到结构((), p1)

    @staticmethod
    def 禁邻位组合(剩余位, 组合, k):
        '''选k个不相邻组合'''
        for pos in 剩余位:
            新组合 = 组合 + [pos]
            if len(新组合) < k:
                BCN.禁邻位组合(set(filter(lambda _: _ > pos, 剩余位 - set(BCN.framework.at[pos]))), 新组合, k)
            else:
                BCN.找到结构([], 新组合)

    @staticmethod
    def 非近邻生成(k):
        BCN.生成初始化()
        BCN.禁邻位组合(set(range(1, len(BCN.framework.at))) - set(BCN.framework.at[0]), [0], k)
        
    @staticmethod
    def 非近邻非五间生成(k):
        def 禁邻位非五间组合(剩余位, 组合, k):
            '''选k个不相邻组合'''
            for pos in 剩余位:
                新组合 = 组合 + [pos]
                for ring in rings5:
                    if len(ring & set(新组合)) >= 2:
                        break
                else:
                    if len(新组合) < k:
                        禁邻位非五间组合(set(filter(lambda _: _ > pos, 剩余位 - set(BCN.framework.at[pos]))), 新组合, k)
                    else:
                        BCN.找到结构([], 新组合)
        
        BCN.生成初始化()
        _, rings5, _ = BCN.骨架ring()
        禁邻位非五间组合(set(range(1, len(BCN.framework.at))) - set(BCN.framework.at[0]), [0], k)

    @staticmethod
    def 等替换生成(k):
        '''替换k个B和k个N'''
        def add_one(自己位, 对方位, 自己候选位):
            for pos in sorted(自己候选位):
                新自己位 = 自己位 + [pos]
                if len(新自己位) != k or len(对方位) != k:
                    对方候选位 = reduce(set.union, [set(BCN.framework.at[_]) for _ in 新自己位]) - \
                               reduce(set.union, [set(BCN.framework.at[_]) for _ in 对方位]) - set(新自己位+对方位)
                    add_one(对方位, 新自己位, 对方候选位)
                else:
                    BCN.找到结构(对方位, 新自己位)

        BCN.生成初始化()
        add_one([], [0], BCN.framework.at[0])

    @staticmethod
    def 限环等替换生成(k):
        '''替换k个B和k个N
        限制有环'''
        def add_one(自己位, 对方位, 自己候选位):
            for pos in sorted(自己候选位):
                新自己位 = 自己位 + [pos]
                if len(新自己位) != k or len(对方位) != k:
                    对方候选位 = reduce(set.union, [set(BCN.framework.at[_]) for _ in 新自己位]) - \
                               reduce(set.union, [set(BCN.framework.at[_]) for _ in 对方位]) - set(新自己位+对方位)
                    add_one(对方位, 新自己位, 对方候选位)
                else:
                    if 交换:
                        BCN.找到结构(新自己位, 对方位)
                    else:
                        BCN.找到结构(对方位, 新自己位)

        BCN.生成初始化()
        B环位 = sorted([0,9,11]) #有序
        N环位 = sorted([1,8,10])
        B候选位 = reduce(set.union, [set(BCN.framework.at[_]) for _ in N环位]) - set(B环位)
        交换 = False
        add_one(B环位, N环位, B候选位)
        N候选位 = reduce(set.union, [set(BCN.framework.at[_]) for _ in B环位]) - set(N环位)
        交换 = True
        add_one(N环位, B环位, N候选位)

    @staticmethod
    def 不等替换生成(n,k):
        '''替换n个B和N，再额外替换k个N'''
        def 禁邻位组合(剩余位, 组合, k):
            '''选k个不相邻组合'''
            for pos in 剩余位:
                新组合 = 组合 + [pos]
                if len(新组合) < k:
                    禁邻位组合(set(filter(lambda _: _ > pos, 剩余位 - set(BCN.framework.at[pos]))), 新组合, k)
                else:
                    BCN.找到结构(B位, 新组合) #这里跟之前代码不同

        BCN.生成初始化()
        path = r'D:\cluster\C44B8N8_ring2'
        for file_name in os.listdir(path):
            相等结构 = read(os.path.join(path,file_name))
            elem_list = 相等结构.get_elements()
            B位 = list(filter(lambda _:elem_list[_]=='B', range(len(elem_list))))
            N位 = list(filter(lambda _:elem_list[_]=='N', range(len(elem_list))))
            N候选位 = reduce(set.union, [set(BCN.framework.at[_]) for _ in B位]) - \
                     reduce(set.union, [set(BCN.framework.at[_]) for _ in N位]) - set(B位+N位)
            禁邻位组合(N候选位, N位, n+k)
            
    def is_五元环间位(self, pos, rings5):
        """rings5: set的列表"""
        for ring in rings5:
            if len(ring & set(pos)) >= 2:
                return True
        return False
            
    def is_六元环3个(self, pos, rings6):
        """rings6: set的列表"""
        for ring in rings6:
            if len(ring & set(pos)) == 3:
                return True
        return False

def BN交换(folder_name, save_folder_name):
    os.mkdir(save_folder_name)
    for file_name in os.listdir(folder_name):
        c = read(os.path.join(folder_name, file_name))
        for a in c:
            if a.elem == 'B':
                a.elem = 'N'
            elif a.elem == 'N':
                a.elem = 'B'
        file_name = file_name.replace('B','X').replace('N','B').replace('X','N')
        write_car(c, os.path.join(save_folder_name, file_name))


#feature16 = ['CCC','CCB','CCN','CBC','CBB','CBN','CNC','CNB','CNN','BCB','BCN','BBB','BBN','BNB','BNN','NCN','NBN','NNN']
feature12 = ['CCC','CCB','CCN','CBC','CBN','CNC','CNB','BCB','BCN','BNB','NCN','NBN']
ring_type = ['CCCCC','BCCCC','BCBCC','NCCCC','NCNCC',
             'CCCCCC','BCCCCC','BCBCCC','BCCBCC','BCBCBC','NCCCCC','NCNCCC','NCCNCC','NCNCNC',
             'CCCBN', 'CCBCN', 'CCBNB', 'BCBNC','CCNBN','NCBNC', 'CBNBN',
             'CCCCBN', 'CCCBCN', 'CCBCCN', 
             'CCCBNB', 'CCBNCB', 'CBNCCB', 'CCCNBN', 'CCNCBN', 'CCBNCN', 'CBCNCN', 'CBCNCB',
             'CCBNBN', 'CBNBCN', 'CBNBCB','CNBNCB','CBNCBN','CBNCNB','CNBNCN', 
             'CNBNBN', 'CBNBNB', 
             'BNBNBN'] #14+28=42

def ring2feature12(ring):
    '''输入ring中元素序号的列表'''
    def permut3_str(permut):
        elem_id = {0:'B', 1:'C', 2:'N'}
        s = reduce(str.__add__, [elem_id[x] for x in permut])
        try:
            return feature12.index(s)
        except:
            return feature12.index(s[::-1])

    n = len(ring)
    ring *= 2
    vector = np.zeros(12, dtype='int')
    for i in range(n):
        vector[permut3_str(ring[i:i+3])] += 1
    return vector
    
    
def ring2int(ring):
    """计算替换环特征"""
    n = len(ring)
    ring *= 2
    vector = np.zeros(12, dtype='int')
    for i in range(n):
        try:
            return ring_type.index(ring[i:i+n])
        except:
            try:
                return ring_type.index(ring[i+n:i:-1])
            except:
                pass


def dual_feature12(id_list, p, edge2face):
    '''输入012的列表，输出对偶图的点特征'''
    feature = np.empty((32,12), dtype='int')
    calc_v = set() #已计算过的对偶点
    for e,v in edge2face.items():
        if v in calc_v:
            continue
        calc_v.add(v)
        ring = p.get_face(e[0],e[1])
        ring = [id_list[_] for _ in ring]
        feature[v] = ring2feature12(ring)
    return feature


def dual_feature_S(id_list, p, edge2face):
    '''输入012的列表，输出对偶图的点特征
    从ring_type里找'''
    elem_id = {0:'B', 1:'C', 2:'N'}
    feature = np.empty((32,), dtype='int')
    calc_v = set() #已计算过的对偶点
    for e,v in edge2face.items():
        if v in calc_v:
            continue
        calc_v.add(v)
        ring = p.get_face(e[0],e[1])
        ring = [id_list[_] for _ in ring]
        s = reduce(str.__add__, [elem_id[x] for x in ring])
        feature[v] = ring2int(s)
    return feature
    

def 最小编号(file_name):
    BCN.g = C60Ih_group()
    c = read(file_name)
    替换 = list(filter(lambda x:c[x].elem=='B', range(len(c))))
    等价替换 = C60Ih轨道(替换, BCN.g)
    print([sorted(B_pos) for B_pos in 等价替换])
    return min([BCN(sorted(B_pos)).get_name() for B_pos in 等价替换])


def 随机选结构(folder_name):
    '''不好用'''
    for file_name in os.listdir(folder_name):
        if hash(file_name) %3 == 0:
            os.remove(os.path.join(folder_name, file_name))
            

def 位置检查(file_name):    
    rings, rings5, rings6 = [], [], []
    backbone = read('C60.car')
    p = Polyhedron.from_cluster(backbone)
    _,edge2face = p.dual()
    calc_v = set() #已计算过的对偶点
    for e,v in edge2face.items():
        if v in calc_v:
            continue
        calc_v.add(v)
        ring = p.get_face(e[0],e[1])
        rings.append(set(ring))
        if len(ring) == 5:
            rings5.append(set(ring))
        if len(ring) == 6:
            rings6.append(set(ring))
    
    f = open(file_name)
    g = open('位置检查.txt','w')
    for line in f:
        pos = [int(_) for _ in line.split('_')[1:]]
        bcn = BCN(pos)
        result = 1 if bcn.is_五元环间位(pos, rings5) else 0
        #result = 1 if bcn.is_六元环3个(pos, rings6) else 0        
        print(result,end='')
        g.write(str(result)+'\n')
            

def test():
    h_list = set()
    for i in range(1,60):
        bcn = BCN([], [0,i])
        h = hash(bcn)
        print(i, h)
        h_list.add(h)
    print(len(h_list))


def test2():
    B_pos, N_pos = [0,11,9,2], [1,10,8,25]
    BCN.g = C60Ih_group()
    等价替换 = C60Ih轨道(B_pos+N_pos, BCN.g)
    for x in 等价替换:
        print(sorted(x[:4]), sorted(x[4:]))
    B_pos2, N_pos2= [0,11,9,7], [1,10,8,4]



if __name__ == '__main__':
    #BCN.单替换(3)
    #BCN.非近邻生成(2)
    #BCN.非近邻非五间生成(11)
    #BCN.等替换生成(8)
    #BCN.限环等替换生成(4)
    #BCN.不等替换生成(8,9)
    #BN交换(r'D:\cluster\C45B7N8_ring2', r'D:\cluster\C45B8N7_ring2')
    #随机选结构(r'D:\cluster\C56N4_forbidden')
    #test()
    #test2()
    #位置检查('pre_data/C53N7_forbidden.txt')
    print(最小编号(r'D:\360disk\Cluster\topology\BCN\BCN_Files\Documents\fine\C53B7-2011 DMol3 GeomOpt\C53B7-2011.car'))